# @Time : 2020/9/4 17:22
# @Author : Fioman 
# @Phone : 13149920693
"""
准备数据集,数据集准备工具
"""
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import json
import os
import cv2 as cv
import numpy as np
import random
from PIL import Image, ImageChops, ImageFilter

trainPath = r"C:\work\code\data_suite\tool2\data\train"
trainJsonPath = r"C:\work\code\data_suite\tool2\data\train\danban.json"
trainData = {}
with open(trainJsonPath, "r") as f:
    trainData = json.load(f)

# 从json文件里面拿到图像文件名和对应的结果,字典对应的键名就是文件名的一个列表集合
trainFilenames = [os.path.join(trainPath, fileName) for fileName in trainData.keys()]
trainTargets = np.array([dictValue["rect"] for dictValue in trainData.values()])

testPath = r"C:\work\code\data_suite\tool2\data\test"
testJsonName = "danban.json"
testJsonPath = os.path.join(testPath, testJsonName)
with open(testJsonPath, "r") as f:
    testData = json.load(f)
testFileNames = [os.path.join(testPath, fileName) for fileName in testData]
testTargets = np.array(dictValue["rect"] for dictValue in testData.values())

# 用于标准化处理,现在没有用
normalize = transforms.Normalize(mean=[0.4], std=[0.09])

# 常规预处理流程,目前只是将图片转换为tensor格式
preprocess = transforms.Compose([transforms.ToTensor(), ])

# 随机预处理,训练时采用这种,给每张图像添加一下亮度对比度饱和度随机变化,提高网络的可靠性
randomPreprocess = transforms.Compose([
    transforms.ColorJitter(brightness=(0.3), contrast=0.1, saturation=0.1, hue=0),
    transforms.ToTensor(),
])


# 默认图像载入函数,数据集工具在提供数据时会自行调用
def default_loader(path):
    imgPil = Image.open(path).convert("L")
    imgTensor = preprocess(imgPil)
    return imgTensor


def image_offset(img, offsetX, offsetY):
    c = ImageChops.offset(img, offsetX, offsetY)
    return c


# 添加了随机变化的图像载入函数,用于提高网络可靠性
def random_loader(path, rotateDegrees=0, offsetX=0, offsetY=0):
    imgPil = Image.open(path).convert("L")
    imgPil = imgPil.ratate(rotateDegrees, resample=False, expand=False, center=None)
